import argparse
import json
import os
import shutil

from utils import resize

def get_template(fp: str) -> dict:
    with open(fp, "r") as f:
        template = f.read()

    return json.loads(template)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--llm-dataset-dir', '-i', type=str, required=True)
    parser.add_argument('--videos-list', '-V', type=str, default='videos.txt')
    parser.add_argument('--prompt-list', '-p', type=str, default='prompt.txt')
    parser.add_argument('--output-dir', '-o', type=str, required=True)
    parser.add_argument('--template', '-t', default='prepare_llm_template.json', type=str)
    parser.add_argument('--json-file', '-j', type=str, default='messages.json')

    args = parser.parse_args()
    args.template = get_template(args.template)

    if os.path.isdir(args.output_dir):
        shutil.rmtree(args.output_dir)

    os.makedirs(os.path.join(args.output_dir, 'videos'), exist_ok=True)

    with open(os.path.join(args.llm_dataset_dir, args.prompt_list), "rt") as f:
        prompts = [i.strip() for i in f.readlines()]

    with open(os.path.join(args.llm_dataset_dir, args.videos_list), "rt") as f:
        videos = [i.strip() for i in f.readlines()]

    messages = []
    for video, prompt in zip(videos, prompts):
        to = f'videos/{os.path.basename(video)}'
        shutil.copy(video, os.path.join(args.output_dir, to), follow_symlinks=False)
        cond = args.template.copy()
        _messages = []
        for message in cond['message']:
            message = message.copy()
            if 'content' in message and isinstance(message['content'], str):
                message['content'] = message['content'].replace('<answer>', prompt)
            _messages.append(message)
        cond['message'] = _messages
        cond['videos'] = [to]
        messages.append(cond)

    with open(os.path.join(args.output_dir, args.json_file), "w") as fp:
        json.dump(messages, fp, indent=2)


if __name__ == '__main__':
    main()